package middleware.redis;

import java.io.IOException;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.Pipeline;
import redis.clients.jedis.Response;
import redis.clients.jedis.Transaction;

/**
 * redis 简单限流：《Redis深度历险：核心原理和应用实践》
 *
 * @author ME
 * @date   2020/3/4
 */
public class SimpleRateLimiter {

    private final Jedis jedis;

    public SimpleRateLimiter(Jedis jedis) {
        this.jedis = jedis;
        // 测试不 connect() 也没异常
        this.jedis.connect();
    }

    /**
     * 简单限流：比如说某用户获取验证码的限流，60s 内 5次
     *
     * @param userId 用户ID
     * @param actionKey 操作类型，如：用户登录验证码操作
     * @param period 窗口大小（单位 /s），如：60s
     * @param maxCount 窗口内的最大次数
     * @return 是否允许操作
     * @throws IOException IO异常
     */
    public boolean isActionAllowed(String userId, String actionKey, int period, int maxCount) throws IOException {
        String key = String.format("hist:%s:%s", userId, actionKey);
        long nowTs = System.currentTimeMillis();

        // 获取管道对象
        Pipeline pipe = jedis.pipelined();
        // 开启事务
        pipe.multi();
        // 以时间为分数，向有序集合添加成员，或更新已有成员的分数
        pipe.zadd(key, nowTs, "" + nowTs);
        // 以时间为分数，删除 1s 之前的所有元素
        pipe.zremrangeByScore(key, 0, nowTs - period * 1000L);
        // 设置过期时间
        pipe.expire(key, period + 1);
        // 获取有序集合的成员数
        // Response<Long> count = pipe.zcard(key);
        Response<Long> count = pipe.zcount(key, 0, nowTs);
        // 执行所有事务块内的命令
        pipe.exec();
        pipe.sync();
        // 释放管道
        pipe.close();

        return count.get() <= maxCount;
    }

    public static void main(String[] args) throws IOException {
        Jedis jedis = new Jedis("localhost");
        SimpleRateLimiter limiter = new SimpleRateLimiter(jedis);
        for (int i = 0; i < 20; i++) {
            System.out.println(limiter.isActionAllowed("laoqian", "reply", 60, 5));
        }
    }
}
